import numpy as np
import matplotlib.pyplot as plt

# 参数设置
Nx, Ny = 32, 32  # 假设我们的矩阵大小为 32x32
Lx, Ly = 1, 1  # 假设我们的域长度为 1x1
dx, dy = Lx / (Nx - 1), Ly / (Ny - 1)  # 计算网格间距

# 波数
kx = 4  # 假设我们在 x 方向有一个特定的波数 kx
ky = 4  # 假设我们在 y 方向有一个特定的波数 ky

# 创建一个二维正弦函数作为输入
x = np.linspace(0, Lx, Nx, endpoint=False)
y = np.linspace(0, Ly, Ny, endpoint=False)
X, Y = np.meshgrid(x, y, indexing='ij')
u = np.sin(2 * np.pi * kx * X) * np.sin(2 * np.pi * ky * Y)

# 使用两次一维FFT模拟二维FFT
u_fft_rows = np.fft.fft(u, axis=0)
u_fft_2d = np.fft.fft(u_fft_rows, axis=1)

# 验证：使用两次一维IFFT重建原始矩阵
u_ifft_rows = np.fft.ifft(u_fft_2d, axis=1)
u_reconstructed = np.fft.ifft(u_ifft_rows, axis=0)

# 计算误差
error = np.linalg.norm(u - u_reconstructed, 'fro')

print("Error after IFFT:", error)

# 显示原始信号与重建信号的差异
diff = np.abs(u - u_reconstructed)

plt.figure(figsize=(15, 5))

plt.subplot(1, 3, 1)
plt.title("Original Signal")
plt.imshow(np.real(u), extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Amplitude')

plt.subplot(1, 3, 2)
plt.title("Reconstructed Signal")
plt.imshow(np.real(u_reconstructed), extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Amplitude')

plt.subplot(1, 3, 3)
plt.title("Difference (Original - Reconstructed)")
plt.imshow(diff, extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Absolute Difference')

plt.tight_layout()
plt.show()

# 计算u_fft_rows的幅度
amp_fft_rows = np.abs(u_fft_rows)

# 计算u_ifft_rows的幅度，以便可视化（尽管通常我们直接比较u与u_reconstructed）
amp_u_ifft_rows = np.abs(u_ifft_rows)

# 绘制u_fft_rows的幅度
plt.figure(figsize=(15, 5))

plt.subplot(1, 2, 1)
plt.title("Amplitude of FFT Rows")
plt.imshow(amp_fft_rows, extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Amplitude')

# 绘制u_ifft_rows的幅度
plt.subplot(1, 2, 2)
plt.title("Amplitude of IFFT Rows (After 1st IFFT)")
plt.imshow(amp_u_ifft_rows, extent=[0, Lx, 0, Ly], origin='lower', cmap='viridis')
plt.colorbar(label='Amplitude')

plt.tight_layout()
plt.show()